from ssummolib import load_index
from dict_to_phyloxml import write_xml
from count_hmms import countseqs
import os, re, sys, time, subprocess
import cPickle as pickle
from CONFIG import *
ambiguous_count = 0
global results          ### This shall be the taxonomic dictionary with all the results in.
                        ### It shall represent tdict, but only includes nodes that have best hits.
                        ### For each node, keep accessions and scores.
                        ### e.g. results[node] = {'accessions' : [acc1, ... ] , 'subnodes' : { ... } }


###############################################################

################ User configured Variables ####################

## These have all been moved to CONFIG.py
## If you need to change these variables, please do so there...

#hmmerdir = '/Users/albl500/bin'
#blastdir = '/Users/albl500/bin'
#blastdbdir = '/Users/albl500/454RDPstuff/blastDB'
#blastversion = 2.2
#top = '/Users/454RDPstuff/arbDB'
#fastacmd = 'blastdbcmd'
hmmsearch = os.path.join(hmmerdir,'hmmsearch')
###############################################################



top = arbDBdir

printHeaders = True

def parse_HMM_pipe(HMM,pipe,nodes):
        global printHeaders
        temp_scores = nodes
        reg = re.compile('\s+')
        result_count = 0
        for line in pipe.stdout:
                if line.startswith('#') or line.strip() == '':
                        if printHeaders:
                                print line.rstrip()
                        continue
                printHeaders = False
                print line.rstrip()
                result_count += 1
                results = reg.split(line)  # Splits the line by whitespace. Iterate through line entries.
                #####################################
                accession = results[0]
                Evalue = float(results[4])
                score = float(results[5])
                #####################################
                if result_count == 1:     ## Only append the HMM to temp_scores['order'] once per pipe...
                        if 'order' in temp_scores.keys():
                                temp_scores['order'].append(HMM)
                        else:
                                temp_scores.update ( { 'order' : [HMM] } )
                else:
                        pass
                if accession in temp_scores.keys():
                        temp_scores[accession][0].append(Evalue)
                        temp_scores[accession][1].append(score)
                else:
                        temp_scores.update({ accession : [ [Evalue], [score] ] } )
        if result_count == 0:
                print '## No matches!'
                return False
        return temp_scores

def parse_hmm(score_dict):
    """This will parse the tabular result file from hmmsearch (must 
    specify --tblout temp_table.txt ) and return the Evalue and
    the score.
    """
#    x = re.compile('\s[0-9\.]+[e\-]*[0-9]*')
    x = re.compile('[a-zA-Z0-9\.\-]+\s+')
    with file('temp_table.txt','r') as handle:
        accessions = set()
        for line in handle:
            score_list = []
            if line.startswith('#'):
                continue
            else:
                for result in x.finditer(line):
                    result = result.group().strip()
                    score_list.append(result)
                accession = score_list[0]
                if '|' in accession:
                    accession = accession[accession.find('|')+1:]
                if accession in accessions:
                    print 'Query sequence {0} appears with scores {1} & {2}. Taking {2}'.format( accession, score_list[1],score_dict[accession][1][-1])
                    continue
                else:
                    accessions.add(accession)
                if accession not in score_dict.keys():
                    score_dict.update( { accession : [ [score_list[4]] ,[ score_list[5] ]] } )
                else:
                    score_dict[accession][0].append( score_list[4] )
                    score_dict[accession][1].append( score_list[5] )
    return score_dict

def diverge( path, accession, equal_choices, tax_node ):
    """Given a path, an accession and multiple best scores, go into each of them
    finding the overall best match"""
    global blastdir
    global blastdbdir
    global fastacmd
    global hmmsearch
    subprocess.call( [ os.path.join( blastdir,fastacmd ), '-d', os.path.join(blastdbdir,'tempDB'), '-i', accession,'-o','temp_seq.fa'], shell=False)
#    os.system('{0}/{1} -d {2}/tempDB -i "{3}" -o temp_seq.fa'.format( blastdir, fastacmd, blastdbdir,  accession))
    temp_dict = {}
    for choice in equal_choices:
        for choices in tax_node[choice].keys():
            subprocess.call([ hmmsearch , '-o', os.devnull, '--tbleout','temp_table.txt','--noali',os.path.join(path,choices,choices+'.hmm'),'temp_seq.fa'],shell=False)
            #os.system( '{0}/hmmsearch -o /dev/null --tblout temp_table.txt --noali "{1}/{2}/{3}.hmm" temp_seq.fa'.format(hmmerdir, path, choices, choices) )
            temp_dict = parse_hmm(temp_dict)
            if x != False:
                accessions_best = score(temp_dict, tax_node[choice])
            else:
                print "no hmm results. Can't do this node. %s/%s" % (path,choice)
                continue
    assert len(accessions_best.keys() ) == 1
    accession, choice = accessions_best.items()[0]
    return {accession : choice}

def score(score_dict, tax_node):
        best_key = None
        tax_node_keys = score_dict.pop('order')
        accessions_best = {}
        Evalue = 10
        score = 1
        num_nodes = len(tax_node_keys)
        diverge_dict = {}
    ##### FIND BEST MATCHES, ASSIGNING ACCESSIONS TO NODES #####
        for accession in score_dict.keys():
                diverge_set = set()
                Evalues = score_dict[accession][0]
                scores = score_dict[accession][1]
                key_len = len(Evalues)
                try:
                        assert key_len == len(scores) and key_len == len(tax_node_keys)
                except AssertionError:
                        print "Problem with accession {0}".format(accession)
                        print "Evalues: {0}".format(Evalues)
                        print "Numbers (lengths):-\n\tkey_len: {0}\tkey_scores: {1}\tnum_nodes: {2}".format(key_len,len(scores),num_nodes)
                        print "Data:\n\tkey_scores: {0}\n\tNodes: {1}".format(scores,tax_node_keys)
                        print "Score_dict items: {0}".format(str([x for x in score_dict.items()]))
                        raise
                best_Evalue = min(Evalues)
                best_score = max(scores)
                n_best = scores.count(best_score)
                if n_best > 1:
                        print '{0} choices for {1}'.format(n_best, accession)
                        prev_index = -1
                        for count in range(n_best):
                                prev_index = scores.index(best_score, prev_index+1)
                                drawing_node = tax_node_keys[prev_index]
                                diverge_set.add(drawing_node)
#                               if drawing_node in diverge_dict.keys():
#                                       diverge_dict[tax_node_keys[prev_index]].append( accession  )
#                               else:
#                                       diverge_set.add( { tax_node_keys[prev_index] : accession } )
                        results = diverge(path,accession,diverge_set,tax_node,hmmsearch)
                        if results == None:
                                print "Can't find a clear match for {0}".format(accession)
                                ambiguous_count += 1
                                results_file.write('{0}\t{1}\n'.format(accession,path) )
                        else:
                                accession, tax_node_key = results
                                accessions_best.update({accession : tax_node_key } )
                elif n_best == 1:
                        print "scores are {0}".format(scores)
                        print "Best score for {0} is {1}".format(accession,best_score)
                        accessions_best.update( { accession : tax_node_keys[scores.index(best_score)] } )
                elif n_best == 0:
                        print 'no matches for accession {0}'.format(accession)
                else:
                        print "n_best = {0}".format(n_best)
        return accessions_best

def fastacmd(accessions_file = '/dev/stdin',entry='-entry_batch'):
        """returns a pipe to blast's fastacmd, now blastdbcmd.
        Can input accessions via returned pipe's .stdin method.
        When extracting all sequences, entry must be changed to
        '-entry' for the recent blast versions."""
        global blastversion
        global blastdbdir
        blastv = int(blastversion.split('.')[-1])
        if blastv < 22:
                fastapipe = subprocess.Popen([ os.path.join(blastdir,'fastacmd'),'-d', os.path.join(blastdbdir,'tempDB'),'-i',accessions_file] ,shell =False,stdin=subprocess.PIPE, stdout = subprocess.PIPE , bufsize=-1)
        elif blastv <= 24:
                #               os.system('{0}/blastdbcmd -db {2}/tempDB -entry_batch accessions.txt -out temp_seqs.fa '.format(blastdir,fastacmd,blastdbdir) )
                fastapipe =  subprocess.Popen([ os.path.join(blastdir,'blastdbcmd'),'-db',os.path.join(blastdbdir,'tempDB'),entry,accessions_file], shell=False, stdin=subprocess.PIPE, stdout=subprocess.PIPE, bufsize=1024)
        else:
                fastapipe =  subprocess.Popen([ os.path.join(blastdir,'blastdbcmd'),'-db',os.path.join(blastdbdir,'tempDB'),entry,accessions_file,'-out','/dev/stdout'], shell=False, stdin=subprocess.PIPE, stdout = subprocess.PIPE ,universal_newlines=True,bufsize=-1)
        return fastapipe

def SSUMMO(node, path,results_node,seqs_file_name='generated'):
        global top
        global hmmerdir
        global hmmsearch
    ## At each path, choose which dir to go into, by performing
    ## hmmsearch on the fasta sequence against each next HMM.
        top_len = len(top)
        node_keys = sorted(node.keys())
        num_nodes = len(node_keys)
        if num_nodes == 0 or node_keys == ['accessions']:
                return results_node
        if path == top and seqs_file_name != 'generated':
                seqs = 'all'
                accessions = []
                with file(seqs_file_name,'rU') as seq_file:
                        sequence_count = 0
                        for line in seq_file:
                                if line.startswith('>'):
                                        accessions.append(line[1:line.find(' ')])
                                        sequence_count += 1
                seqs = '\n'.join(accessions)
        else:
                seqs = '\n'.join(results_node['accessions'])
                sequence_count = len(results_node['accessions'])
        temp_dict = {}  # { accession : [[Evalues, ... ] , [Scores, ... ]] , ... }
    ##### HMMSEARCH & RETRIEVE SCORES #####
        print "## hmmsearching {0} sequences against {1}".format(sequence_count, node_keys)

        for key in node_keys:
                if key == 'accessions':
                        continue
                temp_path = path + '/' + key
                os.mkfifo( 'tempseqs.ind' )
                seqpipe = fastacmd( accessions_file = 'tempseqs.ind')
                with file( 'tempseqs.ind','w' ) as FIFO:
                        FIFO.write( seqs + '\n' )
#                seqpipe.stdin.write(seqs + '\n\n')
#                seqpipe.stdin.close()
                hmmsearchPipe = subprocess.Popen([hmmsearch,'-o',os.devnull,'--tblout','/dev/stdout','--noali',os.path.join(temp_path,key+'.hmm'),'/dev/stdin'], shell=False, stdin=subprocess.PIPE, stdout=subprocess.PIPE,bufsize=-1 )
#                seqs = seqpipe.stdout.read()
                hmmsearchPipe.stdin.write( seqpipe.stdout.read() )
                hmmsearchPipe.stdin.close()
                os.remove( 'tempseqs.ind' )
                temp_dict = parse_HMM_pipe(key,hmmsearchPipe,temp_dict)
                hmmsearchPipe.wait()
                if temp_dict != False:
                        pass
                else:
                        print '## no hmm results. Cant do this node. %s/%s' % (path,key)
                        continue
###         temp_dict.update( { key : [float(Evalue), float(score)] } )
        #del(temp_dict['order'])
        accessions_best = score(temp_dict, node)
    ##### STORE RESULTS / UPDATE RESULT DICT #####
        if accessions_best != {}:
                choices = set()
                for accession, choice in accessions_best.items():
                        if choice not in results_node.keys():
                                results_node.update( { choice : {'accessions' :[accession] } } )
                        else:
                                results_node[choice]['accessions'].append(accession)
                        choices.add(choice)
    ##############################################

                for choice in choices:
                        print "## number of seqs assigned to {0} is {1}".format(repr(choice),len(results_node[choice]['accessions']))
                        ## If node has no results, delete it.
                        if len(results_node[choice]['accessions']) == 0:
                                del(node[choice])
                        ## otherwise recursively do SSuMMO on winning node.
                        else:
                                SSUMMO( node[choice], path + '/' + choice , results_node[choice])
                                del(results_node[choice]['accessions'])
        else:
                print "## No accessions found in {0}".format(choice)
                print node.keys()
                print path
                raise ValueError
                num_nodes = len(nodes)
        return results_node


def prepare(in_sequence_file_name):
        """Given the name of a sequence file, makes a blast database
        and provides the entry point for the SSUMMO algorithm. Returns
        a dictionary of results containing taxonomic nodes with matches
        and at each node the accessions that have been assigned there.
        """
        global results_file
        global top
        global blastversion
        global blastdir
        global blastdbdir
        seq_file_name = in_sequence_file_name
        results_file_name = seq_file_name[:seq_file_name.rfind('.')]
        results_file = file('{0}_results.txt'.format(results_file_name),'w')
        results_file.write('{0}\t{1}\t{2}\t{3}\t{4}\t{5}\n'.format('Name from fasta file', 'Genus guessed from SSUMMO','Seq. Length','Time','Score','Evalue'))
        if blastversion < 2.6:
                subprocess.call([ os.path.join(blastdir,'formatdb') ,'-i',seq_file_name,'-o','T','-n', os.path.join(blastdbdir,'tempDB')], shell=False)
#                       '{0}/formatdb -i "{1}" -o T -n {2}/tempDB'.format(blastdir, seq_file_name, blastdbdir))
        elif blastversion == 2.6:
                subprocess.call([ os.path.join(blastdir,'makeblastdb'),'-in',seq_file_name,'-parse_seqids','-hash_index','-out', os.path.join(blastdbdir,'tempDB')],shell=False)
               # os.system('{0}/makeblastdb -in "{1}" -parse_seqids -hash_index -out {2}/tempDB'.format(blastdir, seq_file_name, blastdbdir) )
        else:
                print "Haven't tested on a blast version this recent. Attempting to use v2.6 config settings"
                subprocess.call([ os.path.join(blastdir,'makeblastdb'),'-in',seq_file_name,'-parse_seqids','-hash_index','-out', os.path.join(blastdbdir,'tempDB')],shell=False)
                #os.system('{0}/makeblastdb -in "{1}" -parse_seqids -hash_index -out {2}/tempDB'.format(blastdir, seq_file_name, blastdbdir) )
        t = time.time()
        results = {}
        print "Formatted blastDB in {0} seconds".format(t - t0)
        try:
                results_dict = SSUMMO( load_index(), top, results,seqs_file_name = seq_file_name)
        except Exception:
                results_file.close()
                raise
        t = time.time() - t0
        print "processed {0} sequences in {1} seconds".format(countseqs(os.getcwd(),file_name = seq_file_name),t)
        return results_dict


if __name__ == '__main__':
        t0 = time.time()
        try:
                results_file_name = sys.argv[1][:sys.argv[1].rfind('.')]
                results_dict = prepare(sys.argv[1])
        except IndexError:
                print "Wrong number of arguments"
                raise
        print "## {0} ambiguous results".format(ambiguous_count)
        ### Save results files
        print "## Saving results"
        with file('{0}.pkl'.format(results_file_name),'wb') as pickle_file:
                print "Saving pickled results file to '{0}.pkl'".format(results_file_name)
                pickle.dump(results_dict,pickle_file)
        with file('{0}.xml'.format(results_file_name),'w') as write_handle:
                print "## Writing phyloxml file to '{0}.xml'".format(results_file_name)
                write_xml(top,write_handle,results_dict)
        ### Create bar graph index file for ITOL ###
        ### CGI over to ITOL ###
        t = time.time()
        mins = int(t-t0) / 60
        print "## Successfully finished in {0} mins{1} secs".format( mins, (t - t0) - (mins*60))





